package com.waterwolf.wolfbot;
import java.util.HashMap;
import java.util.ListIterator;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.FieldNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;

public class MyClassLoader extends ClassLoader implements Opcodes {

    private final JarFile rsJar;

    public HashMap<String, Class<?>> classes = new HashMap<String, Class<?>>();

    public MyClassLoader(final JarFile rsJar) {
        this.rsJar = rsJar;
    }

    @Override
    public Class<?> findClass(final String name) {
        return classes.get(name);
    }

    @Override
    public Class<?> loadClass(final String name) throws ClassNotFoundException {
        if (classes.containsKey(name))
            return classes.get(name);
        Class<?> c = null;
        try {
            c = super.findSystemClass(name);
        } catch (final ClassNotFoundException cnfe) {
            final ZipEntry zip = rsJar.getEntry(name.replaceAll("\\.", "/") + ".class");
            byte[] buffer;
            try {
                if (zip != null) {
                    final ClassReader cr = new ClassReader(rsJar.getInputStream(zip));
                    final ClassNode cn = new ClassNode();
                    cr.accept(cn, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);

                    injectClassLoader(cn);

                    final ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
                    cn.accept(cw);
                    buffer = cw.toByteArray();
                    c = defineClass(name, buffer, 0, buffer.length);
                } else {
                    c = super.loadClass(name);
                }
            } catch (final Exception e) {
                e.printStackTrace();
            }
        }
        classes.put(name, c);
        return c;
    }
   
    public static void injectClassLoader(final ClassNode cn) {
        try {
            if (cn.name.equals("loader")) {
                cn.interfaces.add("com/waterwolf/wolfbot/internal/Loader");
                final ListIterator<FieldNode> fls = cn.fields.listIterator();
                while (fls.hasNext()) {
                    final FieldNode field = fls.next();
                    if (field.desc.equals("Ljava/applet/Applet;")) {
                        final MethodNode node = new MethodNode(17, "getApplet",
                                "()Ljava/applet/Applet;", null, null);
                        node.visitVarInsn(ALOAD, 0);
                        node.visitFieldInsn(GETFIELD, "loader", field.name,
                                "Ljava/applet/Applet;");
                        final int op = Type.getType("Ljava/applet/Applet;")
                                .getOpcode(IRETURN);
                        node.visitInsn(op);
                        node.visitMaxs(op == (DRETURN | LRETURN) ? 2 : 1, 1);
                        node.visitEnd();
                        cn.methods.add(node);
                    }
                }
            }
            if (cn.superName.equals("java/lang/ClassLoader")) {
                
                final String uSupername = "com/waterwolf/wolfbot/ClientClassLoader";
                
                cn.superName = uSupername;
                System.out.println(cn.name + " hacked for custom classloader");
                final ListIterator<?> mli = cn.methods.listIterator();
                while (mli.hasNext()) {
                    final MethodNode mn = (MethodNode) mli.next();
                    if (mn.name.equals("<init>")) {
                        final ListIterator<?> ili = mn.instructions.iterator();
                        while (ili.hasNext()) {
                            final AbstractInsnNode ain = (AbstractInsnNode) ili
                                    .next();
                            if (ain.getOpcode() == Opcodes.INVOKESPECIAL) {
                                final MethodInsnNode min = (MethodInsnNode) ain;
                                min.owner = uSupername;
                                break;
                            }
                        }
                    } else if (mn.name.equals("loadClass")) {
                        final ListIterator<?> ili = mn.instructions.iterator();
                        while (ili.hasNext()) {
                            final AbstractInsnNode ain = (AbstractInsnNode) ili
                                    .next();
                            if (ain.getOpcode() == Opcodes.INVOKEVIRTUAL) {
                                final MethodInsnNode min = (MethodInsnNode) ain;
                                if (min.name.equals("defineClass")) {
                                    min.owner = uSupername;
                                    min.name = "_defineClass";
                                    break;
                                }
                            }
                        }
                    }
                }
            }
        } catch (final Exception e) {
            e.printStackTrace();
        }
    }

}